import random

import gym
from gym.spaces import Discrete

from Games.GameMaze import GameMaze
from Games.MazeDTO import generate_maze


class MazeEnv(gym.Env):
    def __init__(self, minsize=2, maxsize=2):
        self.gameMaze = GameMaze(minsize, maxsize)
        self.size = self.gameMaze.random_maze_size()
        self.Tmaze = generate_maze(self.size, self.size)
        self.Maze = self.Tmaze[0]
        self.maze = self.Maze.cellMaze
        start = self.Tmaze[1]
        end = self.Tmaze[2]
        self.start_cell = self.maze[start[0]][start[1]]
        self.end_cell = self.maze[end[0]][end[1]]
        self.end_cell.is_end = True
        self.cur_pos = self.start_cell
        self.dirs = [[-1, 0], [1, 0], [0, 1], [0, -1]]
        self.action_space = Discrete(4)
        self.observation_space = Discrete(self.size * self.size)
        self.steps = 0

    def reset(self, *args, **kwargs):
        # maze = [[1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1],
        #         [0, 0, 1, 0, 0, 0, 1, 0, 0, 0, 0, 0, 1],
        #         [1, 0, 1, 0, 1, 0, 1, 0, 1, 1, 1, 1, 1],
        #         [1, 0, 0, 0, 1, 0, 1, 0, 1, 0, 1, 0, 1],
        #         [1, 1, 1, 1, 1, 0, 1, 0, 1, 0, 1, 0, 1],
        #         [1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1],
        #         [1, 1, 1, 1, 1, 0, 1, 0, 1, 1, 1, 1, 1],
        #         [1, 0, 1, 0, 0, 0, 1, 0, 1, 0, 0, 0, 1],
        #         [1, 0, 1, 1, 1, 0, 1, 0, 1, 1, 1, 0, 1],
        #         [1, 0, 0, 0, 0, 0, 1, 0, 1, 0, 0, 0, 1],
        #         [1, 0, 1, 0, 1, 1, 1, 0, 1, 0, 1, 0, 1],
        #         [1, 0, 1, 0, 0, 0, 1, 0, 0, 0, 1, 0, 0],
        #         [1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1]]
        #
        # self.Maze.reset_cellMaze(
        #     [[1, 1, 1, 1, 1, 1, 1, 1, 1], [0, 0, 1, 0, 1, 0, 0, 0, 1], [1, 0, 1, 0, 1, 1, 1, 0, 1],
        #      [1, 0, 0, 0, 1, 0, 0, 0, 1], [1, 1, 1, 0, 1, 0, 1, 1, 1], [1, 0, 0, 0, 1, 0, 1, 0, 1],
        #      [1, 0, 1, 1, 1, 0, 1, 0, 1], [1, 0, 0, 0, 0, 0, 0, 0, 0], [1, 1, 1, 1, 1, 1, 1, 1, 1]]
        # )
        self.Maze.reset_cellMaze()
        self.maze = self.Maze.cellMaze
        self.cur_pos = self.start_cell
        loc = self.cur_pos.loc
        self.steps = 0
        first = self.size * loc[0] + loc[1]
        self.cur_pos.visited_num = 1
        tmp = []
        for row in self.maze:
            for cell in row:
                tmp.append(0)
        tmp[first] = 1
        tmp.append(0)
        return tmp

    def step(self, action):
        direction = self.dirs[action]
        pre = self.cur_pos
        loc_pre = pre.loc
        state_pre = loc_pre[0] * self.size + loc_pre[1]
        self.cur_pos = self.gameMaze.move(self.maze, self.cur_pos, self.end_cell, direction)
        loc = self.cur_pos.loc
        state = loc[0] * self.size + loc[1]
        done = self.end_cell.loc == self.cur_pos.loc

        flag = 0
        if done:
            reward = self.size ** 2
        elif state_pre == state:
            flag = 1
            reward = -10
        else:
            reward = -1
        self.steps += 1
        tmp = []
        for row in self.maze:
            for cell in row:
                tmp.append(0)
        tmp[state] = 1
        if self.steps > self.size ** 4:
            done = True
        tmp.append(flag)

        return tmp, reward, done, {}
        # if action == 1:
        #     reward = 1
        # else:
        #     reward = -1
        # if self.steps > self.size ** 2:
        #     done = True
        # self.steps += 1
        # return state, reward, done, {}


if __name__ == '__main__':
    env = MazeEnv(maxsize=3, minsize=3)
    env.reset()
    for i in range(100):
        step = random.randint(0, 3)
        env.step(step)
